#!/usr/bin/env python3
from Bio.PDB import PDBParser
import numpy as np
import argparse
import json
import os
import sys
import pandas as pd
from pathlib import Path


def calculate_box_size(ligand_pdb_file, center_coords, buffer=5.0):
    """
    根据配体尺寸和中心坐标计算对接盒子大小。

    Args:
        ligand_pdb_file (str): 配体的 PDB 文件路径。
        center_coords (tuple): 对接盒子的中心坐标 (x, y, z)。
        buffer (float): 在配体尺寸基础上增加的缓冲距离 (单位：埃)。

    Returns:
        dict: 包含对接盒子中心坐标和尺寸的字典 {  
            "center": (x, y, z),  
            "size": (size_x, size_y, size_z)  
        }  
    """
    try:
        parser = PDBParser()
        structure = parser.get_structure("LIG", ligand_pdb_file)
        min_coords = np.array([float('inf')] * 3)
        max_coords = np.array([float('-inf')] * 3)

        for model in structure:
            for chain in model:
                for residue in chain:
                    for atom in residue:
                        min_coords = np.minimum(min_coords, atom.coord)
                        max_coords = np.maximum(max_coords, atom.coord)

        ligand_size = max_coords - min_coords
        box_size = tuple(ligand_size + 2 * buffer)
        return {"center": center_coords, "size": box_size}
    except Exception as e:
        raise RuntimeError(f"计算盒子尺寸失败: {str(e)}")


def get_center_coords(args):
    """根据优先级获取中心坐标"""
    # 优先级1: 直接指定的xyz坐标
    if all([args.center_x, args.center_y, args.center_z]):
        return (args.center_x, args.center_y, args.center_z)

    # 优先级2: prediction文件
    if args.prediction and os.path.exists(args.prediction):
        try:
            print(f"尝试读取prediction文件: {args.prediction}", file=sys.stderr)
            df = pd.read_csv(args.prediction)
            # df = pd.read_csv("/Pub/Users/wangyk/Project_wangyk/Codelib_YK/codeLibrary_yk/R/Convertional_Transcriptome/Molecular_Docking_cmd/v1.1/test/CTSS_1glo_no_water.pdb_predictions.csv")
            df.rename(columns=lambda x: x.strip().replace(' ', ''), inplace=True)
            required_cols = ['center_x', 'center_y', 'center_z']
            if all(col in df.columns for col in required_cols):
                return tuple(df.loc[0, required_cols])
            else:
                raise ValueError("Prediction文件缺少必要的中心坐标列")
            
        except Exception as e:
            print(f"警告: 读取prediction文件失败: {e}", file=sys.stderr)

    # 默认值 (可根据需求修改)
    print("警告: 未指定中心坐标且prediction文件不可用, 使用默认值(40,40,40)", file=sys.stderr)
    return (40.0, 40.0, 40.0)


def get_output_path(ligand_path, user_specified_path=None):
    """确定输出文件路径"""
    if user_specified_path:
        return user_specified_path
    return f"{Path(ligand_path)}_docking_center_and_autobox_info.json"


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="计算对接盒子参数: 中心坐标优先级: 直接指定 > prediction文件 > 默认值",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument("ligand_pdb", help="配体PDB文件路径")
    parser.add_argument("-x", "--center_x", type=float, help="对接盒子中心X坐标")
    parser.add_argument("-y", "--center_y", type=float, help="对接盒子中心Y坐标")
    parser.add_argument("-z", "--center_z", type=float, help="对接盒子中心Z坐标")
    parser.add_argument("-p", "--prediction", help="包含预测中心坐标的文件路径")
    parser.add_argument("-b", "--buffer", type=float, default=5.0,
                        help="在配体尺寸基础上增加的缓冲距离(埃)")
    parser.add_argument("-o", "--output", help="输出JSON文件路径")

    try:
        args = parser.parse_args()

        # 验证输入文件存在
        if not os.path.exists(args.ligand_pdb):
            raise FileNotFoundError(f"配体PDB文件不存在: {args.ligand_pdb}")

        # 获取中心坐标
        center_coords = get_center_coords(args)

        # 计算盒子参数
        results = calculate_box_size(
            args.ligand_pdb, center_coords, args.buffer)

        # 确定输出路径并保存
        output_file = get_output_path(args.ligand_pdb, args.output)
        with open(output_file, "w") as f:
            json.dump(results, f, indent=4)

        print(f"成功生成对接参数文件: {output_file}")
        # print(json.dumps(results, indent=4))

    except Exception as e:
        print(f"错误: {str(e)}", file=sys.stderr)
        sys.exit(1)
